import tkinter as tk
from tkinter import ttk, scrolledtext, messagebox
import requests
import json
import threading
import base64
from PIL import Image, ImageTk
import io
import time
import asyncio
import aiohttp

# 在文件开头添加 Ollama 相关配置
OLLAMA_API_BASE = "http://localhost:11434"
OLLAMA_MODEL = "deepseek"
AUTO_REPLY_ENABLED = True  # 控制是否启用自动回复

class GewechatGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("Gewechat客户端")
        self.root.geometry("802x914")
        
        # 配置
        self.base_url = "http://127.0.0.1:2531/v2/api"
        self.token = None
        self.appid = None
        self.callback_url = "http://127.0.0.1:8080"  # 基础地址
        self.callback_path = "/v2/api/callback/collect"  # 单独保存回调路径
        
        # 创建GUI组件
        self.auto_reply = tk.BooleanVar(value=AUTO_REPLY_ENABLED)
        self.create_notebook()
        
        # 启动回调服务器
        self.start_callback_server()
        
    def create_notebook(self):
        # 创建标签页
        notebook = ttk.Notebook(self.root)
        notebook.pack(expand=True, fill='both')
        
        # 登录页面
        login_frame = ttk.Frame(notebook)
        notebook.add(login_frame, text='登录')
        self.create_login_frame(login_frame)
        
        # 消息页面
        msg_frame = ttk.Frame(notebook)
        notebook.add(msg_frame, text='消息')
        self.create_message_frame(msg_frame)
        
        # 联系人页面
        contact_frame = ttk.Frame(notebook)
        notebook.add(contact_frame, text='联系人')
        self.create_contact_frame(contact_frame)
        
        # 群聊页面
        group_frame = ttk.Frame(notebook)
        notebook.add(group_frame, text='群聊')
        self.create_group_frame(group_frame)

    def create_login_frame(self, parent):
        # 状态显示
        self.status_label = ttk.Label(parent, text="未登录")
        self.status_label.pack(pady=10)
        
        # 登录按钮
        login_btn = ttk.Button(parent, text="扫码登录", command=self.start_login)
        login_btn.pack(pady=10)
        
        # 二维码显示区域
        self.qr_label = ttk.Label(parent)
        self.qr_label.pack(pady=10)
        
        # 日志显示区域
        self.log_text = scrolledtext.ScrolledText(parent, height=10)
        self.log_text.pack(pady=10, fill='x')
        
    def create_message_frame(self, parent):
        # 消息发送区域
        send_frame = ttk.Frame(parent)
        send_frame.pack(fill='x', pady=10)
        
        # 添加自动回复开关
        auto_reply_check = ttk.Checkbutton(
            send_frame, 
            text="启用AI自动回复", 
            variable=self.auto_reply
        )
        auto_reply_check.pack(side='right', padx=5)
        
        ttk.Label(send_frame, text="接收人:").pack(side='left', padx=5)
        self.to_wxid_entry = ttk.Entry(send_frame)
        self.to_wxid_entry.pack(side='left', padx=5)
        
        ttk.Label(send_frame, text="内容:").pack(side='left', padx=5)
        self.content_entry = ttk.Entry(send_frame)
        self.content_entry.pack(side='left', padx=5, expand=True, fill='x')
        
        send_btn = ttk.Button(send_frame, text="发送", command=self.send_message)
        send_btn.pack(side='left', padx=5)
        
        # 消息显示区域
        self.msg_text = scrolledtext.ScrolledText(parent)
        self.msg_text.pack(expand=True, fill='both')
        
    def create_contact_frame(self, parent):
        # 创建工具栏
        toolbar = ttk.Frame(parent)
        toolbar.pack(fill='x', pady=5)
        
        # 添加按钮
        refresh_btn = ttk.Button(toolbar, text="刷新列表", command=self.refresh_contacts)
        refresh_btn.pack(side='left', padx=5)
        
        search_btn = ttk.Button(toolbar, text="搜索好友", command=self.search_contact)
        search_btn.pack(side='left', padx=5)
        
        add_btn = ttk.Button(toolbar, text="添加好友", command=self.add_contact)
        add_btn.pack(side='left', padx=5)
        
        delete_btn = ttk.Button(toolbar, text="删除好友", command=self.delete_contact)
        delete_btn.pack(side='left', padx=5)
        
        # 联系人列表
        columns = ('wxid', 'nickname', 'remark', 'type')
        self.contact_list = ttk.Treeview(parent, columns=columns, show='headings')
        
        # 设置列标题
        self.contact_list.heading('wxid', text='微信ID')
        self.contact_list.heading('nickname', text='昵称')
        self.contact_list.heading('remark', text='备注')
        self.contact_list.heading('type', text='类型')
        
        # 设置列宽
        self.contact_list.column('wxid', width=150)
        self.contact_list.column('nickname', width=150)
        self.contact_list.column('remark', width=150)
        self.contact_list.column('type', width=100)
        
        # 添加滚动条
        scrollbar = ttk.Scrollbar(parent, orient="vertical", command=self.contact_list.yview)
        self.contact_list.configure(yscrollcommand=scrollbar.set)
        
        # 布局
        self.contact_list.pack(side='left', fill='both', expand=True)
        scrollbar.pack(side='right', fill='y')
        
        # 绑定双击事件
        self.contact_list.bind('<Double-1>', self.on_contact_double_click)

    def create_group_frame(self, parent):
        # 群聊列表
        self.group_list = ttk.Treeview(parent, columns=('wxid', 'name'))
        self.group_list.heading('wxid', text='群ID')
        self.group_list.heading('name', text='群名称')
        self.group_list.pack(expand=True, fill='both')
        
        refresh_btn = ttk.Button(parent, text="刷新群聊", command=self.refresh_groups)
        refresh_btn.pack(pady=10)
    
    def safe_update_gui(self, func):
        """安全地在主线程中更新GUI"""
        if threading.current_thread() is threading.main_thread():
            func()
        else:
            self.root.after(0, func)

    def log(self, message):
        """安全地记录日志"""
        def _log():
            self.log_text.insert('end', f"{message}\n")
            self.log_text.see('end')
        self.safe_update_gui(_log)

    def update_status(self, text):
        """安全地更新状态标签"""
        def _update():
            self.status_label.configure(text=text)
        self.safe_update_gui(_update)

    def update_qr_image(self, photo):
        """安全地更新二维码图片"""
        def _update():
            self.qr_label.configure(image=photo)
            self.qr_label.image = photo
        self.safe_update_gui(_update)

    def start_login(self):
        self.log("开始登录流程...")
        threading.Thread(target=self.login_process).start()
    
    def login_process(self):
        try:
            # 1. 获取token
            response = requests.post(f"{self.base_url}/tools/getTokenId")
            data = response.json()
            
            if data.get('ret') == 200:
                self.token = data.get('data')
                self.log(f"Token获取成功: {self.token}")
                
                # 2. 获取二维码
                headers = {
                    'X-GEWE-TOKEN': self.token,
                    'Content-Type': 'application/json'
                }
                
                qr_response = requests.post(
                    f"{self.base_url}/login/getLoginQrCode", 
                    json={'appId': ''}, 
                    headers=headers
                )
                qr_data = qr_response.json()
                
                if qr_data.get('ret') == 200:
                    qr_img_base64 = qr_data.get('data', {}).get('qrImgBase64')
                    uuid = qr_data.get('data', {}).get('uuid')
                    self.appid = qr_data.get('data', {}).get('appId')
                    
                    # 设置回调地址
                    callback_params = {
                        'token': self.token,
                        'callbackUrl': self.callback_url,  # 使用基础地址
                        'appId': self.appid
                    }
                    self.log(f"设置回调地址参数: {callback_params}")
                    
                    # 先测试回调服务器
                    try:
                        import socket
                        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                        result = sock.connect_ex(('127.0.0.1', 8080))
                        sock.close()
                        
                        if result != 0:
                            self.log("回调服务器未启动")
                            return
                            
                        # 设置回调地址
                        callback_response = requests.post(
                            f"{self.base_url}/tools/setCallback",
                            json=callback_params,
                            headers=headers
                        )
                        callback_data = callback_response.json()
                        self.log(f"回调地址设置响应: {callback_data}")
                    except Exception as e:
                        self.log(f"回调服务器测试失败: {str(e)}")
                        return
                    
                    # 等待回调服务器启动
                    time.sleep(2)
                    
                    # 显示二维码
                    if not all([qr_img_base64, uuid, self.appid]):
                        self.log("获取二维码数据不完整")
                        return
                    
                    try:
                        img_data = base64.b64decode(qr_img_base64.split(',')[1])
                        img = Image.open(io.BytesIO(img_data))
                        photo = ImageTk.PhotoImage(img)
                        self.update_qr_image(photo)
                    except Exception as e:
                        self.log(f"显示二维码出错: {str(e)}")
                        return
                    
                    # 检查登录状态
                    self.check_login_status(uuid)
                else:
                    error_msg = qr_data.get('msg', '未知错误')
                    self.log(f"获取二维码失败: {error_msg}")
            else:
                error_msg = data.get('msg', '未知错误')
                self.log(f"获取Token失败: {error_msg}")
            
        except Exception as e:
            self.log(f"登录过程出错: {str(e)}")
            
    def check_login_status(self, uuid):
        retry_count = 0
        while retry_count < 30:  # 最多尝试30次
            try:
                headers = {'X-GEWE-TOKEN': self.token}
                response = requests.post(f"{self.base_url}/login/checkLogin",
                                      json={'appId': self.appid, 'uuid': uuid},
                                      headers=headers)
                data = response.json()
                
                if data.get('ret') == 200:
                    status = data.get('data', {}).get('status')
                    if status == 2:  # 登录成功
                        self.update_status("已登录")
                        self.log("登录成功！")
                        self.refresh_contacts()
                        self.refresh_groups()
                        break
                    elif status == 3:  # 取消登录
                        self.log("用户取消登录")
                        break
                else:
                    self.log(f"检查状态失败: {data.get('msg', '未知错误')}")
                
                retry_count += 1
                time.sleep(2)
            except Exception as e:
                self.log(f"检查登录状态出错: {str(e)}")
                retry_count += 1
                time.sleep(2)
                
        if retry_count >= 30:
            self.log("登录超时，请重试")
    
    def send_message(self):
        if not self.token or not self.appid:
            messagebox.showerror("错误", "请先登录！")
            return
            
        to_wxid = self.to_wxid_entry.get()
        content = self.content_entry.get()
        
        if not to_wxid or not content:
            messagebox.showerror("错误", "请输入接收人和内容！")
            return
            
        try:
            headers = {
                'X-GEWE-TOKEN': self.token,
                'Content-Type': 'application/json'
            }
            response = requests.post(
                f"{self.base_url}/message/postText",
                json={
                    'appId': self.appid,
                    'toWxid': to_wxid,
                    'content': content
                },
                headers=headers
            )
            data = response.json()
            
            if data.get('ret') == 200:
                self.log(f"发送成功: {content}")
                self.content_entry.delete(0, 'end')
            else:
                self.log(f"发送失败: {data.get('msg', '未知错误')}")
                
        except Exception as e:
            self.log(f"发送消息出错: {str(e)}")
    
    def refresh_contacts(self):
        if not self.token or not self.appid:
            return
            
        try:
            headers = {
                'X-GEWE-TOKEN': self.token,
                'Content-Type': 'application/json'
            }
            response = requests.post(
                f"{self.base_url}/contact/list",
                json={
                    'appId': self.appid
                },
                headers=headers
            )
            data = response.json()
            
            if data.get('ret') == 200:
                contacts = data.get('data', [])
                self.log(f"获取到 {len(contacts)} 个联系人")
                
                # 清空现有列表
                for item in self.contact_list.get_children():
                    self.contact_list.delete(item)
                    
                # 添加新数据
                for contact in contacts:
                    self.contact_list.insert('', 'end', values=(
                        contact.get('wxid', ''),
                        contact.get('nickname', ''),
                        contact.get('remark', ''),
                        self.get_contact_type(contact.get('type', 0))
                    ))
            else:
                self.log(f"获取联系人失败: {data.get('msg', '未知错误')}")
                    
        except Exception as e:
            self.log(f"刷新联系人出错: {str(e)}")

    def get_contact_type(self, type_code):
        """获取联系人类型描述"""
        type_map = {
            0: "未知",
            1: "文件传输助手",
            2: "新的朋友",
            3: "微信团队",
            4: "普通好友",
            5: "群聊"
        }
        return type_map.get(type_code, "未知")

    def search_contact(self):
        """搜索联系人"""
        search_window = tk.Toplevel(self.root)
        search_window.title("搜索联系人")
        search_window.geometry("300x150")
        
        ttk.Label(search_window, text="请输入搜索关键词:").pack(pady=10)
        
        keyword_var = tk.StringVar()
        entry = ttk.Entry(search_window, textvariable=keyword_var)
        entry.pack(pady=10)
        
        def do_search():
            keyword = keyword_var.get().strip()
            if not keyword:
                messagebox.showwarning("提示", "请输入搜索关键词")
                return
                
            try:
                headers = {'X-GEWE-TOKEN': self.token}
                response = requests.post(
                    f"{self.base_url}/contact/search",
                    json={
                        'appId': self.appid,
                        'keyword': keyword
                    },
                    headers=headers
                )
                data = response.json()
                
                if data.get('ret') == 200:
                    contacts = data.get('data', [])
                    if contacts:
                        result = "\n".join([
                            f"微信ID: {c.get('wxid', '')}\n"
                            f"昵称: {c.get('nickname', '')}\n"
                            f"备注: {c.get('remark', '')}\n"
                            for c in contacts
                        ])
                        messagebox.showinfo("搜索结果", result)
                    else:
                        messagebox.showinfo("提示", "未找到匹配的联系人")
                else:
                    messagebox.showerror("错误", f"搜索失败: {data.get('msg', '未知错误')}")
                    
            except Exception as e:
                messagebox.showerror("错误", f"搜索出错: {str(e)}")
            
            search_window.destroy()
        
        ttk.Button(search_window, text="搜索", command=do_search).pack(pady=10)

    def add_contact(self):
        """添加好友"""
        add_window = tk.Toplevel(self.root)
        add_window.title("添加好友")
        add_window.geometry("300x200")
        
        # 创建输入框
        ttk.Label(add_window, text="微信ID:").pack(pady=5)
        wxid_var = tk.StringVar()
        ttk.Entry(add_window, textvariable=wxid_var).pack(pady=5)
        
        ttk.Label(add_window, text="验证消息:").pack(pady=5)
        msg_var = tk.StringVar()
        ttk.Entry(add_window, textvariable=msg_var).pack(pady=5)
        
        def do_add():
            wxid = wxid_var.get().strip()
            msg = msg_var.get().strip()
            
            if not wxid:
                messagebox.showwarning("提示", "请输入微信ID")
                return
                
            try:
                headers = {'X-GEWE-TOKEN': self.token}
                response = requests.post(
                    f"{self.base_url}/contact/add",
                    json={
                        'appId': self.appid,
                        'wxid': wxid,
                        'message': msg
                    },
                    headers=headers
                )
                data = response.json()
                
                if data.get('ret') == 200:
                    messagebox.showinfo("提示", "好友请求已发送")
                else:
                    messagebox.showerror("错误", f"添加好友失败: {data.get('msg', '未知错误')}")
                    
            except Exception as e:
                messagebox.showerror("错误", f"添加好友出错: {str(e)}")
            
            add_window.destroy()
        
        ttk.Button(add_window, text="添加", command=do_add).pack(pady=10)

    def delete_contact(self):
        """删除好友"""
        selection = self.contact_list.selection()
        if not selection:
            messagebox.showwarning("提示", "请先选择要删除的联系人")
            return
            
        item = self.contact_list.item(selection[0])
        wxid = item['values'][0]
        nickname = item['values'][1]
        
        if not messagebox.askyesno("确认", f"确定要删除好友 {nickname} ({wxid}) 吗？"):
            return
            
        try:
            headers = {'X-GEWE-TOKEN': self.token}
            response = requests.post(
                f"{self.base_url}/contact/delete",
                json={
                    'appId': self.appid,
                    'wxid': wxid
                },
                headers=headers
            )
            data = response.json()
            
            if data.get('ret') == 200:
                self.contact_list.delete(selection[0])
                messagebox.showinfo("提示", "好友已删除")
            else:
                messagebox.showerror("错误", f"删除好友失败: {data.get('msg', '未知错误')}")
                
        except Exception as e:
            messagebox.showerror("错误", f"删除好友出错: {str(e)}")

    def on_contact_double_click(self, event):
        """双击联系人时打开聊天窗口"""
        selection = self.contact_list.selection()
        if not selection:
            return
            
        item = self.contact_list.item(selection[0])
        wxid = item['values'][0]
        nickname = item['values'][1]
        
        # 切换到消息页面并填充接收人
        self.root.select(1)  # 切换到消息标签页
        self.to_wxid_entry.delete(0, 'end')
        self.to_wxid_entry.insert(0, wxid)

    def refresh_groups(self):
        if not self.token or not self.appid:
            return
            
        try:
            headers = {
                'X-GEWE-TOKEN': self.token,
                'Content-Type': 'application/json'
            }
            response = requests.post(
                f"{self.base_url}/chatroom/list",
                json={
                    'appId': self.appid
                },
                headers=headers
            )
            data = response.json()
            
            if data.get('ret') == 200:
                groups = data.get('data', [])
                
                # 清空现有列表
                for item in self.group_list.get_children():
                    self.group_list.delete(item)
                    
                # 添加新数据
                for group in groups:
                    self.group_list.insert('', 'end', values=(
                        group.get('wxid', ''),
                        group.get('nickname', '')
                    ))
                self.log(f"成功获取 {len(groups)} 个群聊")
            else:
                self.log(f"获取群聊失败: {data.get('msg', '未知错误')}")
                    
        except Exception as e:
            self.log(f"刷新群聊出错: {str(e)}")

    def start_callback_server(self):
        from http.server import HTTPServer, BaseHTTPRequestHandler
        import json
        
        outer_self = self

        class CallbackHandler(BaseHTTPRequestHandler):
            def do_POST(self):
                try:
                    content_length = int(self.headers['Content-Length'])
                    post_data = self.rfile.read(content_length)
                    data = json.loads(post_data.decode('utf-8'))
                    
                    print(f"\n收到回调请求:")
                    print(f"Path: {self.path}")
                    print(f"Data: {data}")
                    
                    # 处理消息回调
                    if self.path.startswith('/v2/api/callback'):
                        if isinstance(data, dict):
                            outer_self.process_message(data)
                    
                    # 返回成功响应
                    self.send_response(200)
                    self.send_header('Content-Type', 'application/json')
                    self.end_headers()
                    self.wfile.write(json.dumps({
                        'ret': 0,  # 使用 0 表示成功
                        'msg': 'success'
                    }).encode('utf-8'))
                except Exception as e:
                    print(f"处理回调消息出错: {str(e)}")
                    self.send_response(200)  # 即使出错也返回成功
                    self.send_header('Content-Type', 'application/json')
                    self.end_headers()
                    self.wfile.write(json.dumps({
                        'ret': 0,
                        'msg': 'success'
                    }).encode('utf-8'))

            def do_GET(self):
                # 返回成功响应
                self.send_response(200)
                self.send_header('Content-Type', 'application/json')
                self.end_headers()
                self.wfile.write(json.dumps({
                    'ret': 0,
                    'msg': 'success'
                }).encode('utf-8'))

        server = HTTPServer(('127.0.0.1', 8080), CallbackHandler)
        print("回调服务器启动在", self.callback_url)
        threading.Thread(target=server.serve_forever, daemon=True).start()

    async def get_ai_response(self, message):
        """调用 Ollama API 获取 AI 回复"""
        try:
            response = requests.post(
                f"{OLLAMA_API_BASE}/api/generate",
                json={
                    "model": OLLAMA_MODEL,
                    "prompt": message,
                    "stream": False
                }
            )
            
            if response.status_code == 200:
                data = response.json()
                return data.get('response', '抱歉，我暂时无法回复。')
            else:
                print(f"AI服务返回错误: {response.status_code}")
                return "AI 服务暂时不可用"
        except Exception as e:
            print(f"调用 AI 服务出错: {str(e)}")
            return "AI 服务出错"

    def process_message(self, message_data):
        """处理接收到的消息"""
        try:
            print(f"处理消息数据: {message_data}")
            
            # 从回调数据中提取消息内容
            if 'data' in message_data:
                message = message_data['data']
            else:
                message = message_data
                
            msg_type = message.get('type')
            from_wxid = message.get('fromWxid')
            content = message.get('content')
            
            if not all([msg_type, from_wxid, content]):
                print("消息数据不完整")
                return
            
            # 更新消息显示
            def update_msg():
                self.msg_text.insert('end', f"\n{'='*50}")
                self.msg_text.insert('end', f"\n时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
                self.msg_text.insert('end', f"\n发送者: {from_wxid}")
                self.msg_text.insert('end', f"\n内容: {content}")
                self.msg_text.insert('end', f"\n{'='*50}\n")
                self.msg_text.see('end')
            self.safe_update_gui(update_msg)
            
            # 如果启用了自动回复且是文本消息
            if self.auto_reply.get() and msg_type == 1:
                print("准备 AI 回复...")
                
                async def do_reply():
                    try:
                        # 获取 AI 回复
                        ai_response = await self.get_ai_response(content)
                        print(f"AI 回复内容: {ai_response}")
                        
                        # 发送回复
                        headers = {'X-GEWE-TOKEN': self.token}
                        response = requests.post(
                            f"{self.base_url}/message/postText",
                            json={
                                'appId': self.appid,
                                'toWxid': from_wxid,
                                'content': ai_response
                            },
                            headers=headers
                        )
                        data = response.json()
                        print(f"发送回复响应: {data}")
                        
                        if data.get('ret') == 200:
                            def update_reply():
                                self.msg_text.insert('end', f"\n{'='*50}")
                                self.msg_text.insert('end', f"\n时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
                                self.msg_text.insert('end', f"\nAI回复 -> {from_wxid}:")
                                self.msg_text.insert('end', f"\n{ai_response}")
                                self.msg_text.insert('end', f"\n{'='*50}\n")
                                self.msg_text.see('end')
                            self.safe_update_gui(update_reply)
                        else:
                            print(f"发送自动回复失败: {data}")
                            self.log(f"发送自动回复失败: {data.get('msg', '未知错误')}")
                    except Exception as e:
                        print(f"AI回复出错: {str(e)}")
                        self.log(f"AI回复出错: {str(e)}")

                # 在新线程中运行异步函数
                try:
                    asyncio.run_coroutine_threadsafe(do_reply(), self.loop)
                except Exception as e:
                    print(f"启动异步回复出错: {str(e)}")
                    self.log(f"启动异步回复出错: {str(e)}")
                
        except Exception as e:
            print(f"处理消息出错: {str(e)}")
            self.log(f"处理消息出错: {str(e)}")

if __name__ == "__main__":
    # 创建事件循环
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    
    root = tk.Tk()
    app = GewechatGUI(root)
    app.loop = loop  # 保存事件循环引用
    
    # 在新线程中运行事件循环
    def run_loop():
        loop.run_forever()
    
    threading.Thread(target=run_loop, daemon=True).start()
    
    try:
        root.mainloop()
    finally:
        loop.stop()
